﻿using Ndns.ServerHost.Abstractions;
using Ndns.ServerHost.Abstractions.Services;
using System.Threading.Tasks;
using Ndns.Packet.Abstractions;

namespace Ndns.Cache
{
    class DnsCacheMiddleware : IMiddleware
    {
        private ICacheService _cacheService;
        private RequestDelegate _next;
        private ILogger _logger;

        public void RegisterServices(IServiceCollection services, RequestDelegate next)
        {
            _cacheService = services.Get<ICacheService>();
            _next = next;
            _logger = services.Get<ILoggerProvider>().GetLogger<DnsCacheMiddleware>();

            if (_cacheService == null)
            {
                _logger.LogFatal(string.Empty, "No memory cache service is configured!");
            }
        }

        public void Start()
        {
            // nop
        }

        public Task InvokeAsync(DnsContext context)
        {
            if (_cacheService == null)
            {
                return _next(context);
            }

            var request = context.Request;
            if (request.Queries.Count != 1)
            {
                _logger.LogInfo(context.TraceIdentifier, "Can't handle this request because it contains more than one query");

                return _next(context);
            }

            var query = request.Queries[0];
            var cachedAnswers = _cacheService.Get(query);
            if (cachedAnswers != null)
            {
                _logger.LogInfo(context.TraceIdentifier, "Using answers from cache.");

                TransferAnswersToResponse(cachedAnswers, context.Response);
                return Task.CompletedTask;
            }
            
            _logger.LogInfo(context.TraceIdentifier, "No cache entries found for this query.");
            return HandleRequest(context, query);
        }

        private async Task HandleRequest(DnsContext context, DnsQuery query)
        {
            await _next(context);

            var answers = MakeQueryAnswers(context.Response);
            if (answers == null)
            {
                _logger.LogInfo(context.TraceIdentifier, $"No valid answers to be stored to cache.");

                return;
            }

            _logger.LogInfo(context.TraceIdentifier, $"Storing {answers.Count} answer(s) to cache.");
            _cacheService.Set(query, answers, CachePriority.Normal);
        }
        
        private void TransferAnswersToResponse(DnsQueryAnswers answers, DnsResponse response)
        {
            foreach (var item in answers.AnswerRRs)
            {
                response.AnswerRRs.Add(item);
            }
            foreach (var item in answers.AuthoritativeNSs)
            {
                response.AuthoritativeNSs.Add(item);
            }
            foreach (var item in answers.AdditionalRRs)
            {
                response.AdditionalRRs.Add(item);
            }
            answers.ServerAuthoritative = response.ServerAuthoritative;
            answers.RecursionAvailable = response.RecursionAvailable;
            answers.AnswerAuthenticated = response.AnswerAuthenticated;
        }

        private DnsQueryAnswers MakeQueryAnswers(DnsResponse response)
        {
            if (response.ReplyCode != DnsReplyCode.NoError || response.IgnoreRequest == true)
            {
                return null;
            }

            if (response.AnswerRRs.Count == 0 && response.AuthoritativeNSs.Count == 0 && response.AdditionalRRs.Count == 0)
            {
                return null;
            }

            var answers = new DnsQueryAnswers();

            foreach (var item in response.AnswerRRs)
            {
                answers.AnswerRRs.Add(item);
            }
            foreach (var item in response.AuthoritativeNSs)
            {
                answers.AuthoritativeNSs.Add(item);
            }
            foreach (var item in response.AdditionalRRs)
            {
                answers.AdditionalRRs.Add(item);
            }
            response.ServerAuthoritative = answers.ServerAuthoritative;
            response.RecursionAvailable = answers.RecursionAvailable;
            response.AnswerAuthenticated = answers.AnswerAuthenticated;

            return answers;
        }

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                    _cacheService?.Dispose();
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。
                _cacheService = null;

                disposedValue = true;
            }
        }

        // 仅当以上 Dispose(bool disposing) 拥有用于释放未托管资源的代码时才替代终结器。
        // ~DnsCacheMiddleware() {
        //   // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
        //   Dispose(false);
        // }

        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
            // 如果在以上内容中替代了终结器，则取消注释以下行。
            // GC.SuppressFinalize(this);
        }
        #endregion
    }
}
